import numpy as np
import matplotlib.pyplot as plt

def get_positional_encoding(max_len, d_model):
    """
    获取正弦位置编码矩阵
    max_len: 序列的最大长度
    d_model: 嵌入维度
    """
    pos_encoding = np.zeros((max_len, d_model))
    
    for pos in range(max_len):
        for i in range(0, d_model, 2):
            pos_encoding[pos, i] = np.sin(pos / (10000 ** (i / d_model)))
            if i + 1 < d_model:
                pos_encoding[pos, i + 1] = np.cos(pos / (10000 ** ((i + 1) / d_model)))
    
    return pos_encoding

# 参数
max_len = 100  # 序列的最大长度
d_model = 16   # 嵌入维度

# 获取正弦位置编码
pos_encoding = get_positional_encoding(max_len, d_model)

# 绘制位置编码图
print(pos_encoding)
plt.figure(figsize=(10, 8))
plt.pcolormesh(pos_encoding)
plt.xlabel('Embedding Dimension')
plt.ylabel('Position')
plt.title('Sinusoidal Positional Encoding')
plt.colorbar()
plt.show()

